package cn.da.shuai.cool.ai.search.client.controller;

import cn.da.shuai.cool.ai.search.client.common.enums.SSEMsgType;
import cn.da.shuai.cool.ai.search.client.util.SSEServerUtil;
import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

@RestController
@RequestMapping("/rag")
@RequiredArgsConstructor
public class RagController {


    private final VectorStore milvusVectorStore;
    private final ChatClient chatClient;


    @PostMapping("/upload")
    public Object uploadFile(@RequestParam("file") MultipartFile file) {
        String filename = file.getOriginalFilename();
        Resource resource = file.getResource();
        TextReader textReader = new TextReader(resource);
        textReader.getCustomMetadata().put("filename", filename);
        textReader.getCustomMetadata().put("fileSize", file.getSize());
        Optional.ofNullable(filename).ifPresent(v -> textReader.getCustomMetadata().put("fileExtension", v.substring(v.lastIndexOf(".") + 1)));
        textReader.getCustomMetadata().put("fileType", file.getContentType());
        textReader.getCustomMetadata().put("fileCharset", textReader.getCharset().name());
        textReader.getCustomMetadata().put("fileLines", textReader.get().size());
        List<Document> documents = textReader.get();
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> tokenList = splitter.split(documents);
        milvusVectorStore.add(tokenList);
        return true;
    }

    @GetMapping("/search")
    public void search(@RequestParam("query") String query) {
        List<Document> documents = milvusVectorStore.similaritySearch(query);
        String context = "";
        if (CollectionUtil.isNotEmpty(documents)) {
            context = documents.stream().map(Document::getText).collect(Collectors.joining("\n"));
        }
        String ragTemplate = """
                基于以下内容回答问题：
                {context}
                问题：
                {query}
                回答：
                如果不确定就回答不知道，不要根据自己的知识回答，只根据提供的文档内容回答。
                """;
        Prompt prompt = new Prompt(StrUtil.format(ragTemplate, Map.of("context", context, "query", query)));
        Flux<String> contentFlux = chatClient.prompt(prompt).stream().content();
        String content = contentFlux.toStream().peek(r -> SSEServerUtil.send("1", r, SSEMsgType.MESSAGE)).collect(Collectors.joining());
        SSEServerUtil.send("1", content, SSEMsgType.DONE);
    }
}
